#include <ntddk.h>

//UNICODE_STRING deviceName = RTL_CONSTANT_STRING(L"\\Device\\MyDevice");
//UNICODE_STRING symbolicName = RTL_CONSTANT_STRING(L"\\??\\MyDevice");


PCHAR PsGetProcessImageFileName(IN PEPROCESS process);

VOID Unload(PDRIVER_OBJECT pDriverObject)
{
	KdPrint(("%s\n", PsGetProcessImageFileName(PsGetCurrentProcess())));
	UNICODE_STRING symbolicName = RTL_CONSTANT_STRING(L"\\??\\MyDevice");
	IoDeleteSymbolicLink(&symbolicName);	
	IoDeleteDevice(pDriverObject->DeviceObject);
	KdPrint(("Goodbye Driver!\n"));
}

VOID EnumDriver(PDRIVER_OBJECT pDriverObject)
{
	KdPrint(("Driver:%p\n", pDriverObject));
	KdPrint(("Device:%p\n", pDriverObject->DeviceObject));
	KdPrint(("DriverName:%wZ\n", &pDriverObject->DriverName));
	KdPrint(("ServiceName:%wZ\n", &pDriverObject->DriverExtension->ServiceKeyName));
	KdPrint(("Database:%wZ\n", pDriverObject->HardwareDatabase));
}

VOID MemoryTest()
{
	PCHAR pcstr = (PCHAR)ExAllocatePool(PagedPool, 1024);
	if (pcstr == NULL) {
		KdPrint(("memory allocate failed!\n"));
		return STATUS_INSUFFICIENT_RESOURCES;
	}
	KdPrint(("memory allocate succeed!\n"));
	RtlZeroMemory(pcstr, 1024);
	strcpy(pcstr, "this is a memory allocation test!");
	KdPrint(("The buffer:%s\n", pcstr));
}

typedef struct _MYDATA {
	LIST_ENTRY listEntry;
	int number;
}MYDATA, *PMYDATA;

VOID LookasideTest() {
	KdPrint(("Lookaside Test!\n"));
	PMYDATA pData[10];
	PAGED_LOOKASIDE_LIST lookasideList;
	ExInitializePagedLookasideList(&lookasideList, NULL, NULL, 0, sizeof(MYDATA), 'abcd', 0);
	for (int i = 0; i < 10; i++) {
		pData[i] = (PMYDATA)ExAllocateFromPagedLookasideList(&lookasideList);
		pData[i]->number = i;
	}
	for (int i = 0; i < 10; i++) {
		KdPrint(("%d\n", pData[i]->number));
		ExFreeToPagedLookasideList(&lookasideList, pData[i]);
		pData[i] = NULL;
	}
	ExDeletePagedLookasideList(&lookasideList);
}

VOID LinkListTest()
{
	LIST_ENTRY listHead;
	int i = 0;
	InitializeListHead(&listHead);
	for (int i = 0; i < 10; i++) {
		PMYDATA pdata = (PMYDATA)ExAllocatePool(PagedPool, sizeof(MYDATA));
		pdata->number = i;
		InsertHeadList(&listHead, &pdata->listEntry);
	}

	while (!IsListEmpty(&listHead)) {
		PLIST_ENTRY pEntry = RemoveHeadList(&listHead);
		PMYDATA pdata = CONTAINING_RECORD(pEntry, MYDATA, listEntry);
		KdPrint(("%d\n", pdata->number));
		ExFreePool(pdata);
	}
}

VOID RtlMemoryTest()
{
	CHAR buffer[1024] = {0};
	CHAR src[128] = "abc";
	CHAR dst[1024] = "abc";
	const char* testStr1 = "This is a Rtl memory test111!";
	const char* testStr2 = "This is a Rtl memory test222!";
	RtlCopyMemory(buffer, testStr1, strlen(testStr1));
	KdPrint(("%s\n", buffer));
	RtlMoveMemory(buffer, testStr2, strlen(testStr2));
	KdPrint(("%s\n", buffer));
	RtlFillMemory(buffer, 64, 0xcc);
	KdPrint(("%s\n", buffer));
	RtlZeroMemory(buffer, sizeof(buffer));
	if (RtlEqualMemory(src, dst, sizeof(src)))
	{
		KdPrint(("Src is as same as dst."));
	}
}

VOID Tests(PDRIVER_OBJECT pDriverObject)
{
	//EnumDriver(pDriverObject);
	//MemoryTest();
	//LookasideTest();
	//LinkListTest();
	RtlMemoryTest();
}

//#pragma code_seg("PAGE")
//#pragma code_seg("INIT")
NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING pRegPath)
{
	KdPrint(("Hello Driver!\n"));
	KdPrint(("%s\n", PsGetProcessImageFileName(PsGetCurrentProcess())));
	UNICODE_STRING deviceName = RTL_CONSTANT_STRING(L"\\Device\\MyDevice");
	UNICODE_STRING symbolicName = RTL_CONSTANT_STRING(L"\\??\\MyDevice");
	NTSTATUS status = STATUS_SUCCESS;
	pDriverObject->DriverUnload = Unload;
	PDEVICE_OBJECT pDeviceObject;
	status = IoCreateDevice(pDriverObject, 0, &deviceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &pDeviceObject);

	if (!NT_SUCCESS(status)){
		KdPrint(("Device Create Failed!\n"));
		return status;
	}
	pDeviceObject->Flags |= DO_BUFFERED_IO;
	pDeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;
	status = IoCreateSymbolicLink(&symbolicName, &deviceName);
	if (!NT_SUCCESS(status)) {
		KdPrint(("Symbolic Create Failed!\n"));
		IoDeleteDevice(pDeviceObject);
		return status;
	}

	Tests(pDriverObject);
}